import numpy as np
import matplotlib.pyplot as plt


def COA(size, max_iterations, lb, ub, dim):
    # C1 and sigma are used to control the intake of crayfish at different temperatures
    C1 = 0.2
    sigma = 3
    # the temperature most suitable for crayfish
    mu = 25
    # food factor
    C3 = 3

    # Initialize population
    def initialization(size, dim, ub, lb):
        return np.random.uniform(lb, ub, (size, dim))

    # Eq.(4) crayfish intake equation
    def p_obj(temp):
        return (
            C1
            / (np.sqrt(2 * np.pi) * sigma)
            * np.exp(-np.power(temp - mu, 2) / (2 * np.power(sigma, 2)))
        )

    # Example usage:
    # Define a cost function, for example:
    def fobj(x):
        return np.sum(x**2)

    # Define Parameters
    X = initialization(size, dim, ub, lb)
    cuve_f = np.zeros(max_iterations)
    global_Cov = np.zeros(max_iterations)
    X_G = np.zeros(dim)  # optimal position obtained so far
    best_fitness = float("inf")
    fitness_f = np.zeros(size)

    for i in range(size):
        fitness_f[i] = fobj(X[i, :])
        if fitness_f[i] < best_fitness:
            best_fitness = fitness_f[i]
            X_G = X[i, :]

    X_L = X_G.copy()  # the optimal position of the current population
    global_fitness = best_fitness
    cuve_f[0] = best_fitness
    iteration = 0

    while iteration < max_iterations:
        C_2 = 2 - (iteration / max_iterations)  # C2 is a decreasing curve
        temp = np.random.uniform() * 15 + 20  # temperature
        X_shade = (X_G + X_L) / 2  # the cave X_shade for summer vacation
        Xfood = X_G.copy()  # Eq.(10)

        Xnew = np.zeros_like(X)

        for i in range(size):
            if temp > 30:
                if np.random.uniform() < 0.5:
                    # summer resort stage
                    Xnew[i, :] = X[i, :] + C_2 * np.random.uniform(0, 1, dim) * (
                        X_shade - X[i, :]
                    )
                else:
                    # competition stage
                    for j in range(dim):
                        z = np.random.randint(0, size)  # Eq.(9)
                        Xnew[i, j] = X[i, j] - X[z, j] + X_shade[j]  # Eq.(8)
            else:
                # foraging stage
                p = p_obj(temp)
                Q = 3 * np.random.uniform() * fitness_f[i] / fobj(Xfood)  # Food size Q
                if Q > (C3 + 1) / 2:  # The food is too big
                    Xfood = np.exp(-1 / Q) * Xfood  # Eq.(12)
                    for j in range(dim):
                        Xnew[i, j] = X[i, j] + Xfood[j] * p * (
                            np.cos(2 * np.pi * np.random.uniform())
                            - np.sin(2 * np.pi * np.random.uniform())
                        )  # Eq.(13)
                else:
                    Xnew[i, :] = (X[i, :] - Xfood) * p + p * np.random.uniform(
                        0, 1, dim
                    ) * X[
                        i, :
                    ]  # Eq.(14)

        # boundary conditions
        for i in range(size):
            for j in range(dim):
                if np.isscalar(ub):
                    Xnew[i, j] = np.clip(Xnew[i, j], lb, ub)
                else:
                    Xnew[i, j] = np.clip(Xnew[i, j], lb[j], ub[j])

        X_L = Xnew[0, :].copy()
        global_fitness = fobj(X_L)

        for i in range(size):
            # Obtain the optimal solution for the updated population
            new_fitness = fobj(Xnew[i, :])
            if new_fitness < global_fitness:
                global_fitness = new_fitness
                X_L = Xnew[i, :].copy()

            # Update the population to a new location
            if new_fitness < fitness_f[i]:
                fitness_f[i] = new_fitness
                X[i, :] = Xnew[i, :].copy()

                if fitness_f[i] < best_fitness:
                    best_fitness = fitness_f[i]
                    X_G = X[i, :].copy()

        cuve_f[iteration] = best_fitness
        global_Cov[iteration] = global_fitness

        iteration += 1

        if iteration % 50 == 0:
            print(f"COA iter {iteration}/{max_iterations}: {best_fitness}")

    return best_fitness, X_G, cuve_f, global_Cov


# Parameters
# number of search agents
size = 30
# maximum number of iterations
max_iterations = 100
# lb=[lb_1,lb_2,...,lb_n] where lb_n is the lower bound of variable n
# ub=[ub_1,ub_2,...,ub_n] where ub_n is the upper bound of variable n
lb = -100
ub = 100
# number of your variables
dim = 3

best_fun, best_position, cuve_f, global_Cov = COA(size, max_iterations, lb, ub, dim)


# 打印最佳结果
print(f"The best-obtained solution by COA is: {best_position}")
print(f"The best optimal value of the objective function found by COA is: {best_fun}")


# 绘图
fig = plt.figure(figsize=(10, 5))

# 子图1：收敛曲线 - 适应值
ax1 = fig.add_subplot(1, 2, 1)
ax1.semilogy(cuve_f)
ax1.set_title("Best fitness convergence curve")
ax1.set_xlabel("Iteration")
ax1.set_ylabel("Best score obtained so far")

# 子图2：收敛曲线 - 全局协方差
ax2 = fig.add_subplot(1, 2, 2)
ax2.semilogy(global_Cov)
ax2.set_title("X_t fitness curve")
ax2.set_xlabel("Iteration")
ax2.set_ylabel("Local score obtained so far")
ax2.legend(["COA"])

plt.tight_layout()
plt.show()
